#! /usr/bin/env python

"""
Script to concatenate MERRA HDF4 files into a single CSV out file.
"""

import argparse
import csv
import os
import os.path
import pyhdf.SD
import time

def process_merra_file(merra_file, out_file):
    # get the date from the file name
    file_date = merra_file.rsplit('.', 3)[1]
    ds = pyhdf.SD.SD(merra_file, pyhdf.SD.SDC.READ)

    # get the hour, u50m and v50m columns
    hour_seq = [time.strftime('%Y-%m-%d %H:%M:%S',
                              time.strptime('{date} {hour}'.format(
                    date=file_date, hour=int(x)), '%Y%m%d %H'))
                for x in ds.select('time')[0:]]
    u50m_seq = [x for x in ds.select('u50m')[0:, 0, 0]]
    v50m_seq = [x for x in ds.select('v50m')[0:, 0, 0]]

    # append output to csv file
    with open(out_file, 'ab') as csv_file:
        writer = csv.writer(csv_file)
        for row in zip(hour_seq, u50m_seq, v50m_seq):
            writer.writerow(row)
    
def main():
    """
    Main function
    This is where the action begins!
    """
    # set up arg parser
    parser = argparse.ArgumentParser(
        description='Concat multiple MERRA HDF4 SDS into one CSV file with '
        '(time, u50m, v50m) rows')
    parser.add_argument('dir_path', metavar='DIR',
                        help='Directory path containing MERRA HDF4 files '
                        '(can be higher level)')
    parser.add_argument('out_file', metavar='OUT',
                        help='File name for concatenating CSV formatted '
                        'output')
    args = parser.parse_args()

    # remove duplicate files and sort by file name
    merra_files = {}
    for root, dirs, files in os.walk(args.dir_path):
        merra_files.update([(x, os.path.abspath(os.path.join(root, x)))
                            for x in files if x[-4:] == '.hdf'])
    sorted_files = sorted(merra_files.values(), key=os.path.basename)

    # get the full path of out file and truncate it
    out_file = os.path.abspath(args.out_file)
    with open(out_file, 'wb') as out_fo:
        out_fo.truncate()

    # process each MERRA file and append to out file
    for cur_file in sorted_files:
        process_merra_file(cur_file, out_file)

if __name__ == '__main__':
    main()
